import argparse
from math import trunc
import os
import cv2
import sys
import random
import xml.etree.ElementTree as ET

from global_var import getPath

#def process_image(img,min_side):

imgFilePath  = 'digger2\\JPEGImages'
xmlFilePath  = 'digger2\\Annotations'

#计算目标有多大比例进入了新图片
def compute_S_in(rec1,rec2):
    left_max = max(rec1[0],rec2[0])
    right_min = min(rec1[2],rec2[2])
    top_max = max(rec1[1],rec2[1])
    bottom_min = min(rec1[3],rec2[3])

    #两个矩形不相交的情况
    if left_max>=right_min or top_max>=bottom_min:
        return 0
    else: #两个矩形有相交区域的情况
        S1 = (rec1[2]-rec1[0])*(rec1[3]-rec1[1])
        #S2 = (rec2[2]-rec2[0])*(rec2[3]-rec2[1])
        S_cross = (bottom_min-top_max)*(right_min-left_max)
        return S_cross/S1
    return 0

def compute_iou(rec1,rec2):
    left_max = max(rec1[0],rec2[0])
    right_min = min(rec1[2],rec2[2])
    top_max = max(rec1[1],rec2[1])
    bottom_min = min(rec1[3],rec2[3])

    #两个矩形不相交的情况
    if left_max>=right_min or top_max>=bottom_min:
        return 0
    else: #两个矩形有相交区域的情况
        S1 = (rec1[2]-rec1[0])*(rec1[3]-rec1[1])
        S2 = (rec2[2]-rec2[0])*(rec2[3]-rec2[1])
        S_cross = (bottom_min-top_max)*(right_min-left_max)
        return S_cross/(S1+S2-S_cross)

    return 0

def cut(srcImageFile,dstImageFile,srcXmlFile,dstXmlFile,objectList,addpixels=0):

        img = cv2.imread(srcImageFile)
        size = img.shape
        h,w = size[0],size[1]

        #最短边小于416不裁剪
        if min(w,h) <= 416:
            return False

        #判断这个区域有没有目标框
        tree = ET.ElementTree
        tree = ET.parse(os.path.join(srcXmlFile))
        root = tree.getroot()

        #检查所有目标框的最长边，决定裁剪图片的大小
        max_wh = 0
        for obj in root.findall('object'):
            # if obj.find('name').text == "nest1":
            bndbox  = obj.find('bndbox')
            left    = int(bndbox.find('xmin').text)
            top     = int(bndbox.find('ymin').text)
            right   = int(bndbox.find('xmax').text)
            bottom  = int(bndbox.find('ymax').text)
            max_wh = max(max_wh,max(right-left,bottom-top))
        #最长边超过图片的1/2不裁剪
        if max_wh>=h or max_wh>=w:
            return False
            
        pixels = int(max_wh*(1+addpixels))
        pixels = max(pixels,416)
        pixels = min(pixels,h,w)
        #没有检测目标返回不进行裁剪
        if pixels == 0:
            return False

        #根据鸟巢的进入新图片的比例确定新图片的裁剪位置
        confirm_rect = False
        for obj in root.findall('object'):
            if obj.find('name').text in objectList:
                bndbox  = obj.find('bndbox')
                left    = int(bndbox.find('xmin').text)
                top     = int(bndbox.find('ymin').text)
                right   = int(bndbox.find('xmax').text)
                bottom  = int(bndbox.find('ymax').text)
                #计算目标框进入新矩形的比例
                for i in range(1000000//pixels):
                    x = random.randint(0,w-pixels)
                    y = random.randint(0,h-pixels)
                    x2 = x+pixels
                    y2 = y+pixels
                    S_in = compute_S_in((left,top,right,bottom),(x,y,x2,y2))
                    if S_in > 0.8:
                        confirm_rect = True
                        break
            if confirm_rect:
                break

        #未找到有效的目标框
        if not confirm_rect:
            return False
 
        #遍历xml中的检测对象，调整标注框位置
        objects = []        
        for obj in root.findall('object'):
            #判断有没有在裁剪的区域内
            bndbox  = obj.find('bndbox')
            left    = int(bndbox.find('xmin').text)
            top     = int(bndbox.find('ymin').text)
            right   = int(bndbox.find('xmax').text)
            bottom  = int(bndbox.find('ymax').text)
            S_in = compute_S_in((left,top,right,bottom),(x,y,x2,y2))

            if S_in <= 0.1:
                #删除这个对象
                root.remove(obj)
            else:
                #修改box位置,保留这个对象
                left = max(0,left-x)
                top = max(0,top-y)
                right = min(x2,right-x)
                bottom = min(y2,bottom-y)        
                bndbox.find('xmin').text = str(left)
                bndbox.find('ymin').text = str(top)
                bndbox.find('xmax').text = str(right)
                bndbox.find('ymax').text = str(bottom)
                objects.append(obj.find('name').text)

        #保存xml文件和图片文件        
        if len(objects) > 0:
            tree.write(dstXmlFile,encoding="utf-8", xml_declaration=False)
            dstimg = img[y:y2,x:x2]
            cv2.imwrite(dstImageFile,dstimg)
            return True
        else:
            return False

if __name__ == '__main__':
    #前缀列表
    prefixs = ['a_','b_','c_']
    #图片扩大系数
    addpixels = [0.1,0.2,0.3]
    #有效对象列表，不存在有效对象不进行裁剪
    objectList = ['crane','digger','ins_zb','smoke']
    count = 0
    imgPath = os.path.join(getPath(),"JPEGImages")
    xmlPath = os.path.join(getPath(),"Annotations")
    imgFiles = os.listdir(imgPath)
    for fileName in imgFiles:        
        srcImageFile = os.path.join(imgPath,fileName)
        shortName = os.path.splitext(fileName)[0]
        srcXmlFile = os.path.join(xmlPath,shortName+'.xml')
        for p in prefixs:
            dstImageFile = os.path.join(imgPath,p+fileName)
            dstXmlFile = os.path.join(xmlPath,p+shortName+'.xml')
            if cut(srcImageFile,dstImageFile,srcXmlFile,dstXmlFile,objectList,addpixels=addpixels[prefixs.index(p)]):
                count += 1
                print("cut finished:"+p+srcImageFile)
            else:
                print("cut failed:"+p+srcImageFile)
    print('cut count: ',count)




